﻿#include "rb_tree.h"

static RBNodeBase*
RBTreeIncrment(RBNodeBase* x)
{
    if (x->m_right != nullptr)
    {
        return RBNodeBase::minimum(x->m_right);
    }
    else
    {// 向上找到第一个比自己大的祖先
        RBNodeBase* y = x->m_parent;
        while (x == y->m_right)
        {// 跳过比自己小的祖先
            x = y;
            y = y->m_parent;
        }
        if (x->m_right != y)
        {// 非特殊节点header(end)
            x = y;
        }
    }
    return x;
}

static RBNodeBase*
RBTreeDecrement(RBNodeBase *x)
{
    if (x->m_color == RBColor::kRBColorRed
        && x->m_parent->m_parent == x)
    {// 特殊节点header(end)
        x = x->m_right;
    }
    else if (x->m_left != nullptr)
    {// 在子树里找到第一个比自己小的节点
        return RBNodeBase::maximum(x->m_left);
    }
    else
    {// 向上找到第一个比自己小的祖先
        RBNodeBase* y = x->m_parent;
        while (x == y->m_left)
        {// 跳过比自己大的祖先
            x = y;
            y = y->m_parent;
        }
        x = y;
    }
    
    return x;
}
/// 左旋
static void
RBTreeRotateLeft(RBNodeBase* const x, RBNodeBase* &root)
{
    RBNodeBase* const y = x->m_right;
    x->m_right = y->m_left;
    if (y->m_left != nullptr)
    {
        y->m_left->m_parent = x;
    }
    y->m_parent = x->m_parent;

    if (x == root)
    {
        root = y;
    }
    else if (IsLeftChild(x))
    {
        x->m_parent->m_left = y;
    }
    else
    {
        x->m_parent->m_right = y;
    }
    y->m_left = x;
    x->m_parent = y;
}

/// 右旋
static void
RBTreeRotateRight(RBNodeBase* const x, RBNodeBase* &root)
{
    RBNodeBase* const y = x->m_left;
    x->m_left = y->m_right;
    if (y->m_right != nullptr)
    {
        y->m_right->m_parent = x;
    }
    y->m_parent = x->m_parent;

    if (x == root)
    {
        root = y;
    }
    else if (IsLeftChild(x))
    {
        x->m_parent->m_left = y;
    }
    else
    {
        x->m_parent->m_right = y;
    }
    y->m_right = x;
    x->m_parent = y;
}
static void
RBTreeInsertRebalance(RBNodeBase* x, RBNodeBase* &root)
{
    // 新节点的父亲
    RBNodeBase* p = nullptr;
    // 新节点的祖父
    RBNodeBase *g = nullptr;
    // 新节点的叔父
    RBNodeBase* u = nullptr;
    // 新节点必为红色
    while (x != root
            && RBColor::kRBColorRed == x->m_parent->m_color)
    {// 修复双红缺陷
        p = x->m_parent;
        g = p->m_parent;
        u = IsLeftChild(x->m_parent)
                ? x->m_parent->m_parent->m_right
                : x->m_parent->m_parent->m_left;

        if (IsBlack(u))
        {// 叔父节点颜色为黑色或不存在
            bool is_left_x = IsLeftChild(x);
            bool is_left_p = IsLeftChild(p);
            if (is_left_x == is_left_p)
            {// 同侧
                // p成为新的子树树根
                p->m_color = RBColor::kRBColorBlack;
            }
            else
            {// 异侧
                // 2种情况，旋转至同侧后，p成为x的儿子
                //assert(x->m_color == RBColor::kRBColorRed);
                // x成为新的子树树根
                x->m_color = RBColor::kRBColorBlack;
                is_left_x ? RBTreeRotateRight(p, root)
                            : RBTreeRotateLeft(p, root);
            }
            // g必定染成红色
            g->m_color = RBColor::kRBColorRed;
            is_left_p ? RBTreeRotateRight(g, root) 
                        : RBTreeRotateLeft(g, root);
            // 旋转之后，g成为g的儿子
            // 还需要维护root的颜色
            break;
        }
        else
        {// 叔父节点颜色为红色
            p->m_color = RBColor::kRBColorBlack;
            u->m_color = RBColor::kRBColorBlack;
            if (g != root)
            {// 颜色翻转（B树上溢）
                g->m_color = RBColor::kRBColorRed;
            }
            x = g;
        }
    }
    // 根节点永远为黑
    if (root->m_color != RBColor::kRBColorBlack)
    {
        root->m_color = RBColor::kRBColorBlack;
    }
}
/// 红黑树插入算法
static void
RBTreeInsert(const bool insert_left,
                    RBNodeBase* x,
                    RBNodeBase* p,
                    RBNodeBase &header)
{
    x->m_parent = p;
    x->m_left = nullptr;
    x->m_right = nullptr;
    x->m_color = RBColor::kRBColorRed;

    if (insert_left)
    {
        p->m_left = x;
        if (p == &header)
        {// x 是 root
            header.m_parent = x;
            header.m_right = x;
        }
        else if (p == header.m_left)
        {// 维护最小的节点
            header.m_left = x;
        }
    }
    else
    {
        p->m_right = x;
        if (p == header.m_right)
        {// 维护最大的节点
            header.m_right = x;
        }
    }
    RBTreeInsertRebalance(x, header.m_parent);
}

static void
RBTreeEraseRebalance(RBNodeBase* x, RBNodeBase* p, RBNodeBase* &root)
{// 修复双黑缺陷
    // x的兄弟节点
    RBNodeBase* s = nullptr;
    // s的红孩子（若左右孩子皆红，左者优先）
    RBNodeBase* t = nullptr;
    while (IsBlack(x)
            && x != root)
    {
        // TODO
        bool is_left_x = false; //IsLeftChild(x);
        if (p->m_left == x)
        {// 左儿子
            is_left_x = true;
            s = p->m_right;
        }
        else if (p->m_right == x)
        {// 右儿子
            is_left_x = false;
            s = p->m_left;
        }
        else
        {
            if (p->m_left)
            {
                s = p->m_left;
                is_left_x = false;
            }
            else
            {
                s = p->m_right;
                is_left_x = true;
            }
        }
#ifdef DEBUG
        assert(s);
#endif
        if (IsBlack(s))
        {
            t = nullptr;
            bool is_left_t = false;
            if (IsRed(s->m_left))
            {
                is_left_t = true;
                t = s->m_left;
            }
            else if (IsRed(s->m_right))
            {
                is_left_t = false;
                t = s->m_right;
            }
            else
            {
                is_left_t = false;
                t = nullptr;
            }
            if (t)
            {// 黑s有红孩子：BB-1
                //bool is_left_t = IsLeftChild(t);
                if (is_left_x == is_left_t)
                {// 同侧，先做一次旋转
                    is_left_t ? RBTreeRotateRight(s, root)
                              : RBTreeRotateLeft(s, root);
                    // 旋转之后t和s指向的不一样了
                    std::swap(t->m_color, s->m_color);
                    std::swap(t, s);
                }
                is_left_x
                    ? RBTreeRotateLeft(p, root)
                    : RBTreeRotateRight(p, root);
                t->m_color = RBColor::kRBColorBlack;
                s->m_color = p->m_color;
                p->m_color = RBColor::kRBColorBlack;
                break;
            }
            else
            {// 黑s无红孩子：BB-2R或BB-2B
                // 黑s转红
                s->m_color = RBColor::kRBColorRed;
                if (IsRed(p))
                {// BB-2R
                    p->m_color = RBColor::kRBColorBlack;
                    break;
                }
                else
                {// BB-2B
                    x = p;
                    p = x->m_parent;
                }
            }
        }
        else
        {// 红s：BB-3
            is_left_x
                ? RBTreeRotateLeft(p, root)
                : RBTreeRotateRight(p, root);
            std::swap(s->m_color, p->m_color);
            s->m_color = RBColor::kRBColorBlack;
        }
    }
    if (x)
    {
        x->m_color = RBColor::kRBColorBlack;
    }
}

/// 替换父亲
static void
RBTreeTransplant(RBNodeBase* x, RBNodeBase* y, RBNodeBase* &root)
{
    if (x == root)
    {
        root->m_parent->m_parent = y;
        root = y;
    }
    else if (IsLeftChild(x))
    {
        x->m_parent->m_left = y;
    }
    else
    {
        x->m_parent->m_right = y;
    }
    if (y)
    {
        y->m_parent = x->m_parent;
    }
}

/// 红黑树删除算法
static RBNodeBase*
RBTreeErase(RBNodeBase* const z, RBNodeBase &header)
{
    RBNodeBase* &root = header.m_parent;
    // TODO
    RBNodeBase* &left_most = header.m_left;
    RBNodeBase* &right_most = header.m_right;
    RBNodeBase* y = z;
    RBNodeBase* x = nullptr;
    RBNodeBase* x_parent = nullptr;
    RBColor origin_color_y = y->m_color;

    if (z->m_left == nullptr)
    {// 右孩子继承，可能为空
        x = z->m_right;
        x_parent = z->m_parent;
        RBTreeTransplant(z, x, root);
    }
    else if (z->m_right == nullptr)
    {// 左孩子不为空，右孩子为空，左孩子继承
        x = z->m_left;
        x_parent = z->m_parent;
        RBTreeTransplant(z, x, root);
    }
    else
    { // 右子树最小节点继承
        y = RBNodeBase::minimum(z->m_right);
        origin_color_y = y->m_color;
        x = y->m_right;
        x_parent = y;
        if (y->m_parent != z)
        {
            x_parent = y->m_parent;
            RBTreeTransplant(y, x, root);
            y->m_right = z->m_right;
            y->m_right->m_parent = y;
        }
        
        RBTreeTransplant(z, y, root);
        y->m_left = z->m_left;
        y->m_left->m_parent = y;
        y->m_color = z->m_color;
    }
    if (y == z)
    {// 更新最值
        if (left_most == z)
        {
            left_most = z->m_right == nullptr
                            ? z->m_parent
                            : RBNodeBase::minimum(x);
        }
        if (right_most == z)
        {
            right_most = z->m_left == nullptr
                            ? z->m_parent
                            : RBNodeBase::maximum(x);
        }
    }
    {
        if (IsRed(x))
        {// y的一定是黑色
#ifdef DEBUG
            assert(origin_color_y == RBColor::kRBColorBlack);
#endif
            x->m_color = RBColor::kRBColorBlack;
        }
        else if (origin_color_y == RBColor::kRBColorBlack)
        {// 修复双黑缺陷
            RBTreeEraseRebalance(x, x_parent, header.m_parent);
        }
    }
    return z;
}